package com.wang.tool.trace;

import org.slf4j.MDC;

public class SimpleTraceTransmitter implements TraceTransmitter {

    private static final ThreadLocal<TraceContext> threadLocal = new ThreadLocal<>();


    @Override
    public TraceContext get() {
        return threadLocal.get();
    }

    @Override
    public void set(TraceContext context) {
        if (context == null) {
            threadLocal.remove();
            MDC.remove(TraceHolder.TRACE_ID_HEADER);
            MDC.remove(TraceHolder.RPC_ID_HEADER);
        } else {
            threadLocal.set(context);
            MDC.put(TraceHolder.TRACE_ID_HEADER, context.getTraceId());
            MDC.put(TraceHolder.RPC_ID_HEADER, context.getRpcId());
        }
    }
}

